!> author: 左志华
!> date: 2022-06-19
!>
!> Open Source Linalg Library <br>
!> 开源线性代数库 (简易版)
!> @note 使用双精度浮点型
module open_linalg_m

    use, intrinsic :: iso_fortran_env, only: dp => real64, stderr => error_unit
    implicit none
    private

    public :: det, inv, solve, eye, diag, tri, tril, triu

    interface inv
        procedure :: rinv, cinv
    end interface inv

    interface solve
        procedure :: rsolve, csolve
    end interface solve

    interface diag
        procedure :: diag_real_rank2, diag_real_rank1
    end interface diag

contains

    !> calculate inverse of a double precision matrix
    function rinv(a) result(b)
        real(dp), intent(in) :: a(:, :)
        real(dp) :: b(size(a, 2), size(a, 1))
        integer :: ipiv(size(a, 1)), info
        real(dp) :: work(size(a, 2))

        b = a
        ! http://www.netlib.org/lapack/explore-html/d8/ddc/group__real_g_ecomputational_ga8d99c11b94db3d5eac75cac46a0f2e17.html
        call dgetrf(size(a, 1), size(a, 2), b, size(a, 1), ipiv, info)
        if (info < 0) then
            write (stderr, *) 'dgetrf: illegal value in argument ', info
            error stop
        else if (info > 0) then
            write (stderr, *) 'dgetrf: singular matrix, U(i,i) is exactly zero, info = ', info
            error stop
        end if

        ! http://www.netlib.org/lapack/explore-html/d8/ddc/group__real_g_ecomputational_ga1af62182327d0be67b1717db399d7d83.html
        call dgetri(size(a, 2), b, size(a, 1), ipiv, work, size(a, 2), info)
        if (info < 0) then
            write (stderr, *) 'dgetri: illegal value in argument ', info
            error stop
        else if (info > 0) then
            write (stderr, *) 'dgetri: singular matrix, U(i,i) is exactly zero, info = ', info
            error stop
        end if

    end function rinv

    !> calculate inverse of a double precision matrix
    function cinv(a) result(b)
        complex(dp), intent(in) :: a(:, :)
        complex(dp) :: b(size(a, 2), size(a, 1))
        integer :: ipiv(size(a, 1)), info
        complex(dp) :: work(size(a, 2))

        b = a
        ! http://www.netlib.org/lapack/explore-html/d8/ddc/group__real_g_ecomputational_ga8d99c11b94db3d5eac75cac46a0f2e17.html
        call zgetrf(size(a, 1), size(a, 2), b, size(a, 1), ipiv, info)
        if (info < 0) then
            write (stderr, *) 'zgetrf: illegal value in argument ', info
            error stop
        else if (info > 0) then
            write (stderr, *) 'zgetrf: singular matrix, U(i,i) is exactly zero, info = ', info
            error stop
        end if

        ! http://www.netlib.org/lapack/explore-html/d8/ddc/group__real_g_ecomputational_ga1af62182327d0be67b1717db399d7d83.html
        call zgetri(size(a, 2), b, size(a, 1), ipiv, work, size(a, 2), info)
        if (info < 0) then
            write (stderr, *) 'zgetri: illegal value in argument ', info
            error stop
        else if (info > 0) then
            write (stderr, *) 'zgetri: singular matrix, U(i,i) is exactly zero, info = ', info
            error stop
        end if

    end function cinv

    !> solve linear system of double precision
    function rsolve(a, b) result(x)
        real(dp), intent(in) :: a(:, :), b(:, :)
        real(dp) :: x(size(b, 1), size(b, 2))

        real(dp) :: a_(size(a, 1), size(a, 2))
        integer :: ipiv(size(a, 1))
        integer :: info

        a_ = a; x = b
        ! http://www.netlib.org/lapack/explore-html/d7/d3b/group__double_g_esolve_ga5ee879032a8365897c3ba91e3dc8d512.html
        call dgesv(size(a, 1), size(b, 2), a_, size(a, 1), ipiv, x, size(b, 1), info)

        if (info < 0) then
            write (stderr, *) 'dgesv: illegal value in argument ', info
            error stop
        else if (info > 0) then
            write (stderr, *) 'dgesv: singular matrix, U(i,i) is exactly zero, info = ', info
            error stop
        end if

    end function rsolve

    !> solve linear system of double precision
    function csolve(a, b) result(x)
        complex(dp), intent(in) :: a(:, :), b(:, :)
        complex(dp) :: x(size(b, 1), size(b, 2))

        complex(dp) :: a_(size(a, 1), size(a, 2))
        integer :: ipiv(size(a, 1))
        integer :: info

        a_ = a; x = b
        ! http://www.netlib.org/lapack/explore-html/d7/d3b/group__double_g_esolve_ga5ee879032a8365897c3ba91e3dc8d512.html
        call zgesv(size(a, 1), size(b, 2), a_, size(a, 1), ipiv, x, size(b, 1), info)

        if (info < 0) then
            write (stderr, *) 'zgesv: illegal value in argument ', info
            error stop
        else if (info > 0) then
            write (stderr, *) 'zgesv: singular matrix, U(i,i) is exactly zero, info = ', info
            error stop
        end if

    end function csolve

    !> calculate determinant of a double precision matrix
    function det(a) result(d)
        real(dp), intent(in) :: a(:, :)
        real(dp) :: d
        real(dp) :: a_(size(a, 1), size(a, 2))
        integer :: ipiv(size(a, 1)), info, i

        a_ = a
        ! http://www.netlib.org/lapack/explore-html/d8/ddc/group__real_g_ecomputational_ga8d99c11b94db3d5eac75cac46a0f2e17.html
        call dgetrf(size(a, 1), size(a, 2), a_, size(a, 1), ipiv, info)
        if (info < 0) then
            write (stderr, *) 'dgetrf: illegal value in argument ', info
            error stop
        else if (info > 0) then
            write (stderr, *) 'dgetrf: singular matrix, U(i,i) is exactly zero, info = ', info
            error stop
        end if

        d = 1.0_dp
        do i = 1, size(a, 2)
            if (ipiv(i) /= i) then
                d = -d*a_(i, i)
            else
                d = d*a_(i, i)
            end if
        end do

    end function det

    !> constructs the identity matrix <br>
    !> 构建单位矩阵
    pure function eye(m, n) result(a)
        integer, intent(in) :: m            !! number of rows
        integer, intent(in), optional :: n  !! number of columns
        integer, allocatable :: a(:, :)
        integer :: i, j, n_

        if (present(n)) then
            n_ = n
        else
            n_ = m
        end if

        allocate (a(m, n_))

        do concurrent(j=1:n_, i=1:m)
            if (i == j) then
                a(i, j) = 1
            else
                a(i, j) = 0
            end if
        end do

    end function eye

    !> constructs a vector of the diagonal elements of a matrix <br>
    !> 构建矩阵对角元素的向量
    pure function diag_real_rank2(a) result(v)
        real(dp), intent(in) :: a(:, :)
        real(dp) :: v(min(size(a, 1), size(a, 2)))
        integer :: i

        do concurrent(i=1:size(v))
            v(i) = a(i, i)
        end do

    end function diag_real_rank2

    !> constructs a matrix with the diagonal elements of a vector <br>
    !> 构建向量对角元素的矩阵
    pure function diag_real_rank1(v) result(a)
        real(dp), intent(in) :: v(:)
        real(dp) :: a(size(v), size(v))
        integer :: i, j

        do concurrent(j=1:size(v), i=1:size(v))
            if (i == j) then
                a(i, j) = v(i)
            else
                a(i, j) = 0.0_dp
            end if
        end do

    end function diag_real_rank1

    !> returns the lower triangular part of a.
    pure function tril(a, k) result(b)
        real(dp), intent(in) :: a(:, :)
        integer, intent(in), optional :: k
        real(dp) :: b(size(a, 1), size(a, 2))
        integer :: k_, i, j

        if (present(k)) then
            k_ = k
        else
            k_ = 0
        end if

        do concurrent(j=1:size(a, 2), i=1:size(a, 1))
            if (i - k_ < j) then
                b(i, j) = 0.0_dp
            else
                b(i, j) = a(i, j)
            end if
        end do
    end function tril

    !> returns the upper triangular part of a.
    pure function triu(a, k) result(b)
        real(dp), intent(in) :: a(:, :)
        integer, intent(in), optional :: k
        real(dp) :: b(size(a, 1), size(a, 2))
        integer :: k_, i, j

        if (present(k)) then
            k_ = k
        else
            k_ = 0
        end if

        do concurrent(j=1:size(a, 2), i=1:size(a, 1))
            if (i - k_ < j) then
                b(i, j) = a(i, j)
            else
                b(i, j) = 0.0_dp
            end if
        end do
    end function triu

    !> returns a tridiagonal matrix.
    pure function tri(m, n, k) result(a)
        integer, intent(in) :: m
        integer, intent(in), optional :: n, k
        integer, allocatable :: a(:, :)
        integer :: n_, k_, i, j

        if (present(n)) then
            n_ = n
        else
            n_ = m
        end if

        if (present(k)) then
            k_ = k
        else
            k_ = 0
        end if

        allocate (a(m, n_))
        do concurrent(j=1:n_, i=1:m)
            if (i - k_ < j) then
                a(i, j) = 0
            else
                a(i, j) = 1
            end if
        end do

    end function tri

end module open_linalg_m
